# Copyright (c) 2008
#   Matt Pizzimenti (mjpizz@gmail.com, www.livelearncode.com)
#
# All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
# 
# Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
# 
# Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 
# Neither the name of the original author nor the names of contributors
# may be used to endorse or promote products derived from this software
# without specific prior written permission.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import re
import sys
import time
import exceptions
from getpass import getpass
from threading import Thread
from threading import RLock

import utils


class QuitForcedInStreamMonitor(exceptions.Exception): pass
class ExceptionOccurredInStreamMonitor(exceptions.Exception): pass


class IOController(utils.Singleton):
    
    def __init__(self):
        
        def prompt_for_password(prompt):
            self.raw_print("")
            result = getpass(prompt, sys.stderr)
            self.raw_print("")
            return result

        def prompt_for_input(prompt):
            self.raw_print("")
            result = raw_input(prompt)
            self.raw_print("")
            return result

        def raw_print(string):
            sys.stderr.write(string + "\n")
        
        self.prompt_for_password = prompt_for_password
        self.prompt_for_input = prompt_for_input
        self.raw_print = raw_print


class ExpecterFactory(object):
    
    def __init__(self):
        self.responders = []
        self.password_responders = []
        self.auto_responders = {}
        self.auto_password_responders = {}
    
    def add_responder(self, pattern):
        self.responders.append(pattern)
    
    def add_password_responder(self, pattern):
        self.password_responders.append(pattern)
    
    def add_auto_responder(self, pattern, response):
        self.auto_responders[pattern] = response
    
    def add_auto_password_responder(self, pattern, response):
        self.auto_password_responders[pattern] = response
    
    def create_expecter(self, prompt_prefix, stdin, log_callback):
        return Expecter(stdin, prompt_prefix, log_callback, self.responders, self.password_responders, self.auto_responders, self.auto_password_responders)
    
    def remove_responder(self, responder):
        if responder in self.responders:
            self.responders.remove(responder)
        if responder in self.password_responders:
            self.password_responders.remove(responder)
        if responder in self.auto_responders:
            del self.auto_responders[responder]
            

class Expecter(object):
    
    def __init__(self, stdin, prompt_prefix, log_callback, responders = [], password_responders = [], auto_responders = {}, auto_password_responders = {}):
        
        self.stdin = stdin
        self.prompt_prefix = prompt_prefix
        self.log_callback = log_callback

        self.responders = [re.compile(mr) for mr in responders]
        self.password_responders = [re.compile(pr) for pr in password_responders]

        self.auto_responders = {}
        for string_ar_key in auto_responders:
            regex_ar_key = re.compile(string_ar_key)
            self.auto_responders[regex_ar_key] = auto_responders[string_ar_key]

        self.auto_password_responders = {}
        for string_ar_key in auto_password_responders:
            regex_ar_key = re.compile(string_ar_key)
            self.auto_password_responders[regex_ar_key] = auto_password_responders[string_ar_key]
        
        self.string_responders = []
        
        self.last_response = None
    
    def add_string_responder(self, pattern):
        self.string_responders.append(pattern)
    
    def remove_string_responder(self, pattern):
        self.string_responders.remove(pattern)
    
    def log_last_response(self):
        if self.last_response != None:
            self.log_callback(self.last_response)
        else:
            raise ValueError("no response to log")
    
    def respond(self, line):
        
        for auto_responder_regex in self.auto_responders:
            
            if auto_responder_regex.search(line):
                
                auto_responder = self.auto_responders[auto_responder_regex]
                self.stdin.write(auto_responder)
                self.stdin.write("\n")
                self.stdin.flush()
                self.last_response = auto_responder
                return auto_responder

        for password_auto_responder_regex in self.auto_password_responders:

            if password_auto_responder_regex.search(line):

                password_auto_responder = self.auto_password_responders[password_auto_responder_regex]
                self.stdin.write(password_auto_responder)
                self.stdin.write("\n")
                self.stdin.flush()
                self.last_response = "<password hidden>"
                return password_auto_responder

        for password_responder_regex in self.password_responders:

            if password_responder_regex.search(line):

                # TODO: indicate in the prompt that this is a password?
                response_from_user = IOController.get_instance().prompt_for_password("%s%s " % (self.prompt_prefix, line))
                self.stdin.write(response_from_user)
                self.stdin.write("\n")
                self.stdin.flush()
                self.last_response = "<password hidden>"
                return response_from_user
        
        for responder_regex in self.responders:
            
            if responder_regex.search(line):
                
                response_from_user = IOController.get_instance().prompt_for_input("%s%s " % (self.prompt_prefix, line))
                self.stdin.write(response_from_user)
                self.stdin.write("\n")
                self.stdin.flush()
                self.last_response = response_from_user
                return response_from_user
        
        for string_responder in self.string_responders:
            
            if string_responder in line:
                
                response_from_user = IOController.get_instance().prompt_for_input("%s%s " % (self.prompt_prefix, line))
                self.stdin.write(response_from_user)
                self.stdin.write("\n")
                self.stdin.flush()
                self.last_response = response_from_user
                return response_from_user
                
        return None


class StreamMonitorTracker(utils.Singleton):
    
    def __init__(self):
        self.capture_threads = {}
        self.lookup_lock = RLock()
    
    def ask_all_monitors_to_force_quit(self):
        self.lookup_lock.acquire()
        try:
            
            if len(self.capture_threads.keys()) > 0:
            
                # collect all pending prompts
                possible_prompts = {}
                for stream_monitor in self.capture_threads:
                    possible_prompt = stream_monitor.get_possible_prompt()
                    if possible_prompt:
                        possible_prompts[stream_monitor] = possible_prompt
            
                if len(possible_prompts.values()) > 0:

                    # we have possible prompts
                    IOController.get_instance().raw_print("Pending stream data:")
                    for stream_monitor in possible_prompts:
                        possible_prompt = possible_prompts[stream_monitor]
                        IOController.get_instance().raw_print("\t" + possible_prompt)
                
                    # ask the user if they want to try responding
                    response = IOController.get_instance().prompt_for_input("Process may be waiting for input. [r]espond, [i]gnore, or [k]ill? ")

                    if response == "r" or response == "respond":
                        # user wanted to try responding
                        for stream_monitor in possible_prompts:
                            possible_prompt = possible_prompts[stream_monitor]
                            stream_monitor.respond_to_prompt(possible_prompt)
                    
                    elif response == "k" or response == "kill":
                        # user wanted to kill everything
                        for stream_monitor in self.capture_threads:
                            stream_monitor.force_quit()

                    else:
                        # user wanted to ignore
                        pass
            
                else:
                
                    # we have no possible prompts
                    response = IOController.get_instance().prompt_for_input("Process may be hung. [i]gnore or [k]ill? ")
                
                    if response == "k" or response == "kill":
                        # user wanted to kill everything
                        for stream_monitor in self.capture_threads:
                            stream_monitor.force_quit()
                    else:
                        # user wanted to ignore
                        pass
            else:
                raise SystemExit()
                
        finally:
            self.lookup_lock.release()

    def start_and_track(self, stream_monitor):
        self.lookup_lock.acquire()
        try:
            if stream_monitor in self.capture_threads:
                raise RuntimeError("monitor already running")
            new_capture_thread = Thread( target = stream_monitor.capture, name = stream_monitor.description )
            self.capture_threads[stream_monitor] = new_capture_thread
            new_capture_thread.start()
        finally:
            self.lookup_lock.release()
    
    def join(self, stream_monitor):
        thread_to_join = None
        self.lookup_lock.acquire()
        try:
            if stream_monitor not in self.capture_threads:
                raise RuntimeError("monitor was not running")
            thread_to_join = self.capture_threads[stream_monitor]
        finally:
            self.lookup_lock.release()
        
        # join it and remove from the pool
        thread_to_join.join()
        self.lookup_lock.acquire()
        try:
            del self.capture_threads[stream_monitor]
        finally:
            self.lookup_lock.release()


class StreamMonitor(object):
    """Helper class that can be used to monitor the data coming
    across a stream and respond to input prompts using the given
    Expecter.  All data on the stream is logged via +log_callback+"""

    description = property(lambda self: self._description)
    output = property(lambda self: self._combined_stream_data)

    def __init__(self, stream, stream_log_callback, expecter = None, description = None, force_quit_callback = None):
        
        self._stream = stream
        self._stream_log_callback = stream_log_callback
        self._expecter = expecter
        if not self._expecter:
            raise ValueError("expecter cannot be 'None'")
        self._description = description
        self._force_quit_callback = force_quit_callback
        if not self._force_quit_callback:
            raise ValueError("force_quit_callback cannot be 'None'")
        
        self._combined_stream_data = ""
        self._waiting_for_stream_data = None
        self._unlogged_stream_data = ""
        self._capturing = False
        self._forced_to_quit = False
        self._unexpected_exception = None

        self._lock = RLock()
    
    def _purge_unlogged_stream_data(self):
        """purges any unlogged stream data, adding it to the logs
        and to the final output.  NOTE: this is not a threadsafe method"""
        self._combined_stream_data += self._unlogged_stream_data
        self._combined_stream_data += "\n"
        self._stream_log_callback(self._unlogged_stream_data)
        self._unlogged_stream_data = ""

    def start(self):
        """Starts the monitoring thread"""
        StreamMonitorTracker.get_instance().start_and_track(self)
        
    def join(self):
        """Joins the internal threads"""
        StreamMonitorTracker.get_instance().join(self)
        if self._forced_to_quit:
            raise QuitForcedInStreamMonitor("manually forced a stream monitor quit")
        if self._unexpected_exception:
            raise ExceptionOccurredInStreamMonitor("an unexpected exception occurred in the StreamMonitor (%s)" % self._unexpected_exception[1])
    
    def force_quit(self):        
        
        # force the quit
        self._force_quit_callback()
        
        # set internal quit state
        self._lock.acquire()
        try:
            self._forced_to_quit = True
        finally:
            self._lock.release()
    
    def get_possible_prompt(self):
        """returns a prompt, if there is any possible prompt waiting on the stream"""
        possible_prompt = None
        self._lock.acquire()
        try:
            if self._capturing and self._waiting_for_stream_data and len(self._unlogged_stream_data.strip()) > 0:
                possible_prompt = self._unlogged_stream_data
        finally:
            self._lock.release()
        return possible_prompt
    
    def respond_to_prompt(self, possible_prompt):
        """given a prompt, tries to respond to it"""
        self._expecter.add_string_responder(possible_prompt)
        if self._expecter.respond(possible_prompt):
            self._purge_unlogged_stream_data()
            self._expecter.log_last_response()
        else:
            raise StandardError("the expecter failed to respond for some unknown reason")
        self._expecter.remove_string_responder(possible_prompt)
        
    def capture(self):
        """The main thread method.  Captures data 1 byte at a time
        and logs it for every newline.  Input is expect-ed properly and
        logged AFTER the prompt that triggered it."""
        
        try:
            
            self._lock.acquire()
            try:
                # update internal state to indicate that we are capturing
                self._capturing = True
                self._waiting_for_stream_data = True
            finally:
                self._lock.release()
        
            # read initial data
            new_data = self._stream.read(1)
        
            # loop until we don't get any new data, or until forcibly stopped
            while new_data != "" and not self._forced_to_quit:
            
                self._lock.acquire()
                try:
                
                    # we have stream data
                    self._waiting_for_stream_data = False
                
                    # check to see if we finished a line of output
                    if new_data == "\n":
                    
                        self._purge_unlogged_stream_data()

                    else:
                    
                        # did not get a whole new line, just concatenate
                        self._unlogged_stream_data = "%s%s" % (self._unlogged_stream_data, new_data)

                        # try responding to the current data as if it were a prompt
                        if self._expecter and self._expecter.respond(self._unlogged_stream_data):
                        
                            # responded to a matched pattern, consider this a "line"
                            self._purge_unlogged_stream_data()
                            self._expecter.log_last_response()
                
                    # reset the state of recent data acquisition
                    # before we unlock and try to read data from the stream
                    self._waiting_for_stream_data = True
                
                finally:
                    self._lock.release()
                
                # read another byte of data
                new_data = self._stream.read(1)
            
            # cleanup
            self._lock.acquire()
            try:
                self._waiting_for_stream_data = False
                self._capturing = False
                if len(self._unlogged_stream_data) > 0:
                    self._purge_unlogged_stream_data()
            finally:
                self._lock.release()
        
        except:
            # track the unhandled exception
            self._unexpected_exception = sys.exc_info()
            
            # cleanup
            self._lock.acquire()
            try:
                self._waiting_for_stream_data = False
                self._capturing = False
            finally:
                self._lock.release()
